import glob
import json
import os
import random

import cv2

'''
labelmepose（模型生成，再手动调整的数据） 转 ClsImgDataset
'''


def labelmepose2yoloClsImgDataset(call_config):
    # img_glob = r"D:\DATA\20250519RENBAO\caitu\test10\*.jpg"
    # # save_dir = r"D:\DATA\20250519RENBAO\caitu\test10_pose_cls_format_data"  # poseimg
    # save_dir = r"D:\DATA\20250519RENBAO\caitu\cls_format_data"  # poseimg

    img_glob = call_config.get('img_glob')
    save_dir = call_config.get('save_dir')
    augment_repeat = call_config.get('augment_repeat', 1)
    is_draw = call_config.get('is_draw', True)
    target_cls_names = call_config.get('target_cls_names')

    ls_img = glob.glob(img_glob)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    train_save_dir = rf"{save_dir}\train"
    test_save_dir = rf"{save_dir}\test"

    # todo 统计不同label个数，利用augment_repeat使个数一致
    # 统计所有 label 出现次数
    from collections import defaultdict
    label_counter = defaultdict(int)  # str:int

    for img_path in ls_img:
        le = len(img_path.split('.')[-1])
        json_path = img_path[:-le] + 'json'
        if not os.path.exists(json_path):
            continue
        with open(json_path, 'r') as f:
            json_data = json.load(f)
        for shape in json_data['shapes']:
            if shape['shape_type'] == 'rectangle':
                cls_name = shape['label']
                # if cls_name not in ['person_Other', 'Other']:
                if cls_name in target_cls_names:
                    label_counter[cls_name] += 1

    print("Label counts:", dict(label_counter))

    # max_count = max(label_counter.values())
    mean_count = sum(label_counter.values()) // len(label_counter)
    cls_augment_repeat_dict = {}  # str:int  cls_name: repeat
    for cls_name in label_counter.keys():
        if cls_name in ['BMFangBao', 'BMZhanLi','AnJianYuan']:
            cls_augment_repeat_dict[cls_name] = 1
        else:
            cls_augment_repeat_dict[cls_name] = max(1, augment_repeat * mean_count // label_counter[cls_name])
    print("cls_augment_repeat_dict:", cls_augment_repeat_dict)

    for img_ind, i in enumerate(ls_img):  # per img

        try:
            img = cv2.imread(i)
            le = len(json_path.split('.')[-1])
            json_path = i[:-le] + '.json'
            with open(json_path, 'r') as f:
                json_data = json.load(f)

            # to crop_infor_list
            # augment_repeat = 5
            # augment_repeat = 1
            crop_infor_list = []  # {'xyxy':xyxy, 'label': str, kp_name1: [x,y] ,,,}
            augment_ls = ['drop', 'noise']
            is_drop = 'drop' in augment_ls
            is_noise = 'noise' in augment_ls

            for shape_dict in json_data['shapes']:  # per target
                if not shape_dict['shape_type'] == 'rectangle':
                    continue

                cls_name = shape_dict['label']
                [[x1, y1], [x2, y2]] = shape_dict['points']
                if cls_name in ['person_Other', 'Other']:
                    continue

                for cls_augment_count in range(cls_augment_repeat_dict[cls_name]):  # 增强次数
                    crop_infor = {}
                    crop_infor['xyxy'] = [x1, y1, x2, y2]
                    crop_infor['label'] = cls_name
                    crop_infor['augment_infor'] = ''
                    crop_infor['augment_count'] = cls_augment_count
                    crop_infor['group_id'] = shape_dict['group_id']
                    group_id = shape_dict['group_id']
                    for kp_shape_dict in json_data['shapes']:  # todo 提速
                        if not (kp_shape_dict['group_id'] == group_id and kp_shape_dict['shape_type'] == 'point'):
                            continue

                        kp_name = kp_shape_dict['label']
                        kpx, kpy = kp_shape_dict['points'][0][0], kp_shape_dict['points'][0][1]
                        if cls_augment_count == 0:
                            crop_infor[kp_name] = [kpx, kpy]
                        else:
                            if is_drop and random.random() > 0.9:
                                kpx, kpy = 0, 0
                                crop_infor['augment_infor'] += f'Drop_{kp_name} '
                            elif is_noise and random.random() > 0.6:
                                ruler_x = int((x2 - x1) / 20)
                                ruler_y = int((y2 - y1) / 20)
                                kpx, kpy = kpx + random.randint(-ruler_x, ruler_x), kpy + random.randint(-ruler_y,
                                                                                                         ruler_y)
                                crop_infor['augment_infor'] += f'Noise_{kp_name} '
                            crop_infor[kp_name] = [kpx, kpy]

                    crop_infor_list.append(crop_infor)

            # crop img
            for crop_ind, crop_infor_dict in enumerate(crop_infor_list):

                # crop
                # img_crop = to_pose_img(img, crop_infor_dict, is_zero_img=True)
                img_crop = to_pose_img(img, crop_infor_dict, is_zero_img=False, is_draw=is_draw)
                cls_name = crop_infor_dict['label']
                augment_infor = crop_infor_dict['augment_infor']
                augment_count = crop_infor_dict['augment_count']
                group_id = crop_infor_dict['group_id']
                # 保存
                train_label_save_dir = f'{train_save_dir}/{cls_name}'  # cls_format_data/train/cls1/
                train_label_save_path = f'{train_label_save_dir}/{cls_name}_img{img_ind}_group{group_id}_aug{augment_count}_{augment_infor}.jpg'
                if not os.path.exists(train_label_save_dir):
                    os.makedirs(train_label_save_dir)
                cv2.imwrite(train_label_save_path, img_crop)  # 训练集

                if random.random() < 0.4:  # 40% 测试集
                    test_label_save_dir = f'{test_save_dir}/{cls_name}'
                    test_label_save_path = f'{test_label_save_dir}/{cls_name}_img{img_ind}_group{group_id}_aug{augment_count}_{augment_infor}.jpg'
                    if not os.path.exists(test_label_save_dir):
                        os.makedirs(test_label_save_dir)
                    cv2.imwrite(test_label_save_path, img_crop)  # 测试集
                print(f'{img_ind}/{len(ls_img)}  {crop_ind}/{len(crop_infor_list)} {cls_name} {save_dir} {i}')


        except:
            print(f'error {i}')

    print("Label counts:", dict(label_counter))
    print("cls_augment_repeat_dict:", cls_augment_repeat_dict)
    print(f'save_dir {save_dir}')

    # train_save_dir
    # ls = glob.glob(f'{save_dir}/train/*/*.jpg')


def to_pose_img(img, crop_infor_dict, is_zero_img=False, is_draw=True):
    kp_names = ['left_shoulder', 'right_shoulder', 'left_elbow', 'right_elbow', 'left_wrist', 'right_wrist']
    kp_names_colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (0, 0, 0), (255, 255, 255), (255, 255, 0)]
    lines = [('left_shoulder', 'left_elbow'), ('left_elbow', 'left_wrist'), ('right_shoulder', 'right_elbow'),
             ('right_elbow', 'right_wrist')]  # 显示用
    lines_colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 255)]
    xyxy = list(map(int, crop_infor_dict['xyxy']))

    # 加pad
    x1, y1, x2, y2 = xyxy
    # pad_x = int((x2 - x1) * 0.5)
    # pad_y = int((y2 - y1) * 0.5)
    pad_xl = int((x2 - x1) * 0.5)
    pad_xr = int((x2 - x1) * 0.5)
    pad_yu = int((y2 - y1) * 0.5)
    pad_yb = int((y2 - y1) * 0.5)

    # pad_xl = int((x2 - x1) * 1)
    # pad_xr = int((x2 - x1) * 1)
    # pad_yu = int((y2 - y1) * 0.25)
    # pad_yb = int((y2 - y1) * 1.25)

    p_x1, p_y1 = max(0, x1 - pad_xl), max(0, y1 - pad_yu)
    p_x2, p_y2 = min(img.shape[1], x2 + pad_xr), min(img.shape[0], y2 + pad_yb)



    # img_crop = img[p_y1:p_y2, p_x1:p_x2, ...]  # crop
    # pad后 再调整为方形
    offset = int(abs((p_x2 - p_x1) - (p_y2 - p_y1)) / 2)
    if p_x2 - p_x1 < p_y2 - p_y1:
        pad2_xl = offset
        pad2_xr = offset
        pad2_yu = 0
        pad2_yb = 0
    else:
        pad2_xl = 0
        pad2_xr = 0
        pad2_yu = offset
        pad2_yb = offset
    p2_x1, p2_y1 = max(0, p_x1 - pad2_xl), max(0, p_y1 - pad2_yu)
    p2_x2, p2_y2 = min(img.shape[1], p_x2 + pad2_xr), min(img.shape[0], p_y2 + pad2_yb)

    pad_black_l = 0
    pad_black_u = 0
    pad_black_r = 0
    pad_black_b = 0
    if p_x1 - pad2_xl < 0:
        pad_black_l = - (p_x1 - pad2_xl )
    if p_y1 - pad2_yu < 0:
        pad_black_u = - (p_y1 - pad2_yu)
    if p_x2 + pad2_xr > img.shape[1]:
        pad_black_r = (p_x2 + pad2_xr) - img.shape[1]
    if p_y2 + pad2_yb > img.shape[0]:
        pad_black_b = (p_y2 + pad2_yb) - img.shape[0] # 边缘时，填充方形

    img_crop = img[p2_y1:p2_y2, p2_x1:p2_x2, ...]  # crop
    img_crop = img_crop.copy()
    if is_zero_img:
        img_crop[:, :, :] = 0

    if not is_draw:
        return img_crop
    # line_width = max(2, int((img_crop.shape[0] + img_crop.shape[1]) * 0.001))
    line_width = max(4, int((img_crop.shape[0] + img_crop.shape[1]) * 0.01))
    # 画线
    for line, line_color in zip(lines, lines_colors):
        if crop_infor_dict.get(line[0]) and crop_infor_dict.get(line[1]):
            kx1, ky1 = list(map(int, crop_infor_dict[line[0]]))
            kx2, ky2 = list(map(int, crop_infor_dict[line[1]]))
            if (kx1 == 0 and ky1 == 0) or (kx2 == 0 and ky2 == 0):
                continue
            # kx1 = kx1 - p_x1
            # ky1 = ky1 - p_y1
            # kx2 = kx2 - p_x1
            # ky2 = ky2 - p_y1
            kx1 = kx1 - p2_x1
            ky1 = ky1 - p2_y1
            kx2 = kx2 - p2_x1
            ky2 = ky2 - p2_y1
            cv2.line(img_crop, (kx1, ky1), (kx2, ky2), line_color, line_width)
            # print()
    # 画点
    for kp_name, kp_name_color in zip(kp_names, kp_names_colors):
        if crop_infor_dict.get(kp_name):
            kx, ky = list(map(int, crop_infor_dict[kp_name]))
            if kx == 0 and ky == 0:
                continue
            # kx = kx - p_x1
            # ky = ky - p_y1
            kx = kx - p2_x1
            ky = ky - p2_y1
            cv2.circle(img_crop, (kx, ky), line_width, kp_name_color, -1)
            # todo 添加噪声
    # cv2.imwrite(fr'D:\DATA\20250519RENBAO\temp\crop_img_{xyxy[0]}_{xyxy[1]}.jpg', img_crop)

    # 人物在边缘 填充为方形
    is_pad_square= False
    if is_pad_square:
        img_crop = cv2.copyMakeBorder(img_crop, top=pad_black_u,bottom=pad_black_b,left=pad_black_l,right=pad_black_r,borderType=cv2.BORDER_CONSTANT,value=(0,0,0))
    return img_crop


if __name__ == '__main__':
    # img_glob = r"D:\DATA\20250519RENBAO\caitu\20250527_annotated_100_200\*.jpg"
    # # img_glob = r"D:\DATA\20250519RENBAO\trainV8Pose_people\add_imgs\*\*.jpg"
    # # img_glob = r"D:\DATA\20250519RENBAO\trainV8Pose_people\add_imgs\*\000001.jpg"
    # # save_dir = r"D:\DATA\20250519RENBAO\caitu\test10_pose_cls_format_data"  # poseimg
    # # save_dir = r"D:\DATA\20250519RENBAO\caitu\cls_format_data_img100_200_cls7_aug5"  # poseimg
    # # save_dir = r"D:\DATA\20250519RENBAO\caitu\cls_format_data_img100_200_cls7_augMean3_padSquare"  # poseimg
    # save_dir = r"D:\DATA\20250519RENBAO\caitu\cls_format_data_img100_200_cls7__augMean3_padSquare_nodraw"  #


    img_glob = r'D:\DATA\20250519RENBAO\trainV8Pose_closePeople\add_imgs\*\*.jpg'
    save_dir = r'D:\DATA\20250519RENBAO\trainV8Pose_closePeople\cls_format_data__labelmepose400_cls4_augMean3_padSquare'
    call_config = {
        'img_glob': img_glob,
        'save_dir': save_dir,
        'augment_repeat': 3,
        'is_draw': True,
        'target_cls_names': ['ZhCMFangBao', 'ZhCMFuKuang',
                          'ZhCMNaKuang', 'ZhCMZhanLi'] ,
    }
    labelmepose2yoloClsImgDataset(call_config)
